# =============================================================== #
# Reproducibility of:                                             #
# Drivers and Trends for the Equality of Opportunity for Sexual   #
# and Gender Minorities: A Panel Approach Equality of Opportunity #
# for Sexual and Gender Minorities 2024                           #
#                                                                 #
# Code written by Omar Alburqueque and reviewed by Paola Ballon   #
# Contact: oalburquequechav@worldbank.org, pballon@worldbank.org  #
# =============================================================== #

# -------------------- #
# Table 4 and Figure 2 #
# -------------------- #

library(readxl)
library(tidyverse)
library(cluster)
library(e1071)
library(clusterCrit)
library(writexl)
library(ggplot2)

project_root    <- getwd()
intermediate_dir <- file.path(project_root, "intermediate files")
outputs_dir      <- file.path(project_root, "outputs")

# Load dataset
data_long <- read_excel(file.path(intermediate_dir, "eqosogi_score.xlsx"))

# Build feature set robustly
features_df <- data_long %>%
  group_by(country) %>%
  arrange(year) %>%
  nest() %>%
  mutate(
    slope_total = map_dbl(data, ~ coef(lm(es_ ~ year, data = .x))[2]),  # Total slope
    delta = map_dbl(data, ~ last(.x$es_) - first(.x$es_)),              # Total change
    # first_value = map_dbl(data, ~ first(.x$es_)),                       # Initial value
    last_value = map_dbl(data, ~ last(.x$es_)),                         # Final value
    sd = map_dbl(data, ~ sd(.x$es_)),                                   # Standard deviation
    mean = map_dbl(data, ~ mean(.x$es_)),                               # Mean
    
    max_jump = map_dbl(data, ~ max(diff(.x$es_))),                      # Largest increase
    # year_of_max_jump = map_dbl(data, ~ .x$year[which.max(diff(.x$es_)) + 1]),  # Year of largest jump
    # min_jump = map_dbl(data, ~ min(diff(.x$es_))),                      # Largest drop
    # n_neg_jumps = map_dbl(data, ~ sum(diff(.x$es_) < 0)),               # Number of declines
    # prop_zero_jumps = map_dbl(data, ~ mean(diff(.x$es_) == 0)),        # Share with no change
    # 
    # slope_early = map_dbl(data, ~ coef(lm(es_ ~ year, data = filter(.x, year >= 1960 & year < 1981)))[2]),  # Early-period slope
    # slope_mid   = map_dbl(data, ~ coef(lm(es_ ~ year, data = filter(.x, year >= 1981 & year < 2001)))[2]),  # Mid-period slope
    slope_late  = map_dbl(data, ~ coef(lm(es_ ~ year, data = filter(.x, year >= 2001)))[2]),                  # Late-period slope
    # 
    # delta_mid   = map_dbl(data, ~ .x$es_[.x$year == 2000] - .x$es_[.x$year == 1981]),  # Mid-period change
    delta_late  = map_dbl(data, ~ .x$es_[.x$year == 2024] - .x$es_[.x$year == 2001]),   # Late-period change
    # 
    # iqr = map_dbl(data, ~ IQR(.x$es_))  # Interquartile range
  ) %>%
  select(-data) %>%
  ungroup()

# Scale features (excluding country)
features_scaled <- scale(select(features_df, -country))

# Apply k-means clustering
set.seed(123)
kmeans_result <- kmeans(features_scaled, centers = 3)
features_df$cluster <- kmeans_result$cluster

# Join with long data for plotting
data_plot <- data_long %>%
  inner_join(features_df %>% select(country, cluster), by = "country")

# Reassign cluster names and order based on visual inspection criteria
data_plot <- data_plot %>%
  mutate(cluster = case_when(
    cluster == 1 ~ "Cluster 1",  # flat
    cluster == 2 ~ "Cluster 3",  # steepest growth
    cluster == 3 ~ "Cluster 2"   # intermediate
  )) %>%
  mutate(cluster = factor(cluster, levels = c("Cluster 1", "Cluster 2", "Cluster 3")))

data_plot <- data_plot %>%
  mutate(es_ = es_ * 100)

# ---------------------------------------------------------- #
# Figure 2: EQOSOGI score trajectories by cluster, 1961–2024 #
# ---------------------------------------------------------- #

# Polished visualization
figure2 <- ggplot(data_plot, aes(x = year, y = es_, group = country)) + 
  geom_line(alpha = 0.4, color = "steelblue") +
  stat_summary(aes(group = cluster), fun = mean, geom = "line", color = "red", size = 0.9) +
  facet_wrap(~ cluster, labeller = label_value, strip.position = "bottom") +
  labs(
    title = "",
    x = "",
    y = "EQOSOGI Score"
  ) +
  theme_minimal() +
  theme(
    plot.title = element_text(hjust = 0.5),      # Center title
    panel.grid.major = element_blank(),          # Remove major gridlines
    panel.grid.minor = element_blank(),          # Remove minor gridlines
    strip.placement = "outside",                 # Place facet labels below
    strip.background = element_blank(),          # Remove facet label background
    strip.text = element_text(face = "bold")     # Bold facet labels
  )

ggsave(
  filename = paste0(outputs_dir, "/Figure2.png"),
  plot     = figure2,
  width    = 8,
  height   = 5,
  dpi      = 300
)

# --------------------------------------------- #
# Table 4: Clustering results for set 4 and k=3 #
# --------------------------------------------- #

# Create collapsed table
cluster_export <- data_plot %>%
  distinct(country, cluster) %>%
  arrange(cluster, country)

write_xlsx(cluster_export, paste0(intermediate_dir, "/eqosogi_cluster_k3.xlsx"))

df <- cluster_export |>
  mutate(cluster = factor(cluster,
                          levels = c("Cluster 1", "Cluster 2", "Cluster 3")))

tabla_clusters <- df |>
  group_by(cluster) |>
  summarise(
    n_countries = n(),
    Countries   = paste(country, collapse = ", "),
    .groups = "drop"
  ) |>
  mutate(
    Cluster = case_when(
      cluster == "Cluster 1" ~
        paste0("Cluster 1: Limited progress, reversals, and low inclusion (",
               n_countries, " countries)"),
      cluster == "Cluster 2" ~
        paste0("Cluster 2: Gradual improvement and moderate inclusion (",
               n_countries, " countries)"),
      cluster == "Cluster 3" ~
        paste0("Cluster 3: Accelerated progress and high inclusion (",
               n_countries, " countries)")
    )
  ) |>
  select(Cluster, Countries)

# Export to Excel
write_xlsx(tabla_clusters, paste0(outputs_dir, "/Table4.xlsx"))
